import scipy.stats as stats
import numpy as np
from scipy.stats import pearson3

class CalculationUnit:
    def __init__(self, uid, area, runoff_coefficient, timet ,rainfall=None, Rainfall95=None, Rainfall75=None, Rainfall50=None):
        self.uid = uid #编号 字符串
        self.area = area #面积 float
        self.runoff_coefficient = runoff_coefficient #径流系数 float
        self.timet = timet #月还是旬 字符串TEN_DAY 和 MONTH
        self.rainfall = rainfall    #降雨量 字典 {key是整型年份，value是有12个或者36个数的列表}
        self.Rainfall95 = Rainfall95  #95降雨量，12个或者36个数的列表
        self.Rainfall75 = Rainfall75  #同上
        self.Rainfall50 = Rainfall50  #同上


        self.year95 = None  #95对应的年份 整形
        self.year75 = None     #同上
        self.year50 = None     #同上
        self.MonRunoffPre95 = []    #95对应的月径流量 列表
        self.MonRunoffPre75 = []
        self.MonRunoffPre50 = []
        self.YearRunoffPre95 = None     #95对应的年径流量值 float
        self.YearRunoffPre75 = None
        self.YearRunoffPre50 = None
    #计算皮尔逊三星曲线参数
    def calculate_parameter(self,data):
        location = np.mean(data)
        scale = np.std(data)
        shape = stats.skew(data)
        return location,scale,shape
    # X: 洪水量或者降水量
    # EX: 均值
    # CV: 离势系数
    # CS: 偏度系数
    def hy_pearson3_x2p(self,x: float, location: float, scale: float, shape: float):
        # return pearson3.sf = 1 - pearson3.cdf
        return pearson3.sf(x, shape, loc=location, scale= scale)
    #通过降雨量求三个水平年降雨
    def calculate_frequency_rainfall(self):
        #根据每年的降雨总量求频率
        sorted_years = sorted(self.rainfall.keys(), key=lambda year: sum(self.rainfall[year]), reverse=True)
        rainfall_value = list(self.rainfall.values())
        rainfall_sorted=[]
        for value in rainfall_value:
            rainfall_sorted.append(sum(value))
        rainfall_sorted=sorted(rainfall_sorted,reverse=True)
        location,scale,shape = self.calculate_parameter(rainfall_sorted)
        sorted_indices = []
        for rain in rainfall_sorted:
            pp = self.hy_pearson3_x2p(rain, location, scale, shape)
            sorted_indices.append(pp)
        #找到距离95，75，50最近的年份序号
        near_95 = min(sorted_indices, key=lambda x: abs(x - 0.95))
        near_75 = min(sorted_indices, key=lambda x: abs(x - 0.75))
        near_50 = min(sorted_indices, key=lambda x: abs(x - 0.5))
        #找到95，75，50对应的年份
        self.year95 = sorted_years[sorted_indices.index(near_95)]
        self.year75 = sorted_years[sorted_indices.index(near_75)]
        self.year50 = sorted_years[sorted_indices.index(near_50)]
        #将三个年份的降雨量分别存入Rainfall95，Rainfall75，Rainfall50当中用于后续计算
        self.Rainfall95 = self.rainfall[self.year95]
        self.Rainfall75 = self.rainfall[self.year75]
        self.Rainfall50 = self.rainfall[self.year50]

    def calculate_mon_runoff_pre(self):
        if self.timet == 'TEN-DAY':
            # 每三旬的降雨相加，旬数据转换成月数据
            monthly_rainfall95 = [sum(self.Rainfall95[i:i+3]) for i in range(0, len(self.Rainfall95), 3)]
            monthly_rainfall75 = [sum(self.Rainfall75[i:i+3]) for i in range(0, len(self.Rainfall75), 3)]
            monthly_rainfall50 = [sum(self.Rainfall50[i:i+3]) for i in range(0, len(self.Rainfall50), 3)]
        else:
            monthly_rainfall95 = self.Rainfall95
            monthly_rainfall75 = self.Rainfall75
            monthly_rainfall50 = self.Rainfall50
        # 通过计算得到月径流量
        for i in range(1, 13):
            runoff_pre95 = self.runoff_coefficient * self.area * monthly_rainfall95[i - 1] * 0.1
            runoff_pre75 = self.runoff_coefficient * self.area * monthly_rainfall75[i - 1] * 0.1
            runoff_pre50 = self.runoff_coefficient * self.area * monthly_rainfall50[i - 1] * 0.1
            self.MonRunoffPre95.append(runoff_pre95)
            self.MonRunoffPre75.append(runoff_pre75)
            self.MonRunoffPre50.append(runoff_pre50)
    #月径流量求和得到地区年径流量
    def calculate_year_runoff_pre(self):
        self.YearRunoffPre95 = sum(self.MonRunoffPre95)
        self.YearRunoffPre75 = sum(self.MonRunoffPre75)
        self.YearRunoffPre50 = sum(self.MonRunoffPre50)










